"""
-*- coding: utf-8 -*-

@author: Du Changping
@time: 2021/12/10 14:39
@file name: create_input_file
@software：PyCharm

Do not smash your computer!

"""

import os
import time
import argparse
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt


def path_setting(model_name_: str):
    """

    :param model_name_: model name to create the folder
    :return: current dir, the folder path to save the vox file
    """
    cur_dir_ = os.path.dirname(__file__)
    save_output_file_dir_path_ = os.path.join(cur_dir_, "output", model_name_)
    if not os.path.exists(os.path.join(cur_dir_, "output")):
        os.mkdir(os.path.join(cur_dir_, "output"))
    if not os.path.exists(save_output_file_dir_path_):
        os.mkdir(save_output_file_dir_path_)
    return cur_dir_, save_output_file_dir_path_


def add_timestamp():
    time_stamp_ = time.strftime("%m%d-%H%M", time.localtime())
    return time_stamp_


def create_in_file(scale_: float,
                   model_name_: str,
                   time_stamp_="",
                   vox_file_path_=None,
                   photon_num_="30820037",
                   gpu_="3",
                   cal_dose_=False,
                   dose_scale_="1"):
    cur_dir_, _ = path_setting(model_name_)
    save_in_file_dir_path_ = os.path.join(cur_dir_, "input", model_name_)
    if not os.path.exists(os.path.join(cur_dir_, "input")):
        os.mkdir(os.path.join(cur_dir_, "input"))
    if not os.path.exists(save_in_file_dir_path_):
        os.mkdir(save_in_file_dir_path_)
    save_in_file_path_ = os.path.join(save_in_file_dir_path_, model_name_ + time_stamp_ + "_{}.in".format(scale_))
    # log_file_path_ = os.path.join(save_in_file_dir_path_, model_name_ + time_stamp_ + ".txt")

    # basic setting todo
    photon_num_ = int(float(photon_num_) * float(dose_scale_))  # TOTAL NUMBER OF HISTORIES, OR SIMULATION TIME IN SECONDS IF VALUE < 100000
    print("photon_num_ is {}".format(photon_num_))
    random_seed_ = 1024  # RANDOM SEED
    gpu_no_ = gpu_  # GPU NUMBER TO USE WHEN MPI IS NOT USED, OR TO BE AVOIDED IN MPI RUNS
    thread_num_ = 512  # GPU THREADS PER CUDA BLOCK (multiple of 32)
    simulated_num_per_thread_ = 1000

    # image setting todo
    nv_ = int(474)  # detector number
    nu_ = int(648)
    du_ = 0.0616  # detector distance (cm)
    dv_ = 0.0616

    # X-ray source setting todo
    # as for the position of the source,
    # we set the it above the center and assume the source , the centre of
    # object and the center of detector are in the one line
    kvp_path_ = os.path.join(cur_dir_, "energy", "liuhanxi_task.spc")
    sdd_ = 125  # source to detector distance(cm)
    source_to_ratation_axis_distance_ = 75
    # sod = 950000  # source to object distance (cm)

    # object size setting todo
    nx = int(512)  # object dimension on x
    ny = int(512)  # object dimension on y
    nz = int(249)  # object dimension on z

    d_x = 0.5 * 0.1  # convert to unit: cm
    d_y = 0.5 * 0.1
    d_z = 1 * 0.1

    source_position_x_ = nx * d_x * scale_ / 2  # X, Y, Z(cm)
    source_position_z_ = nz * d_z * scale_ / 2  # X, Y, Z(cm)  todo still need to consider the position
    source_position_y_ = ny * d_y * scale_ / 2 - source_to_ratation_axis_distance_

    # scan trajectory setting todo
    projection_num_ = 180
    angle_between_projection_ = 1
    angle_of_interest_ = [270, 270 + 3600]
    # source_to_ratation_axis_distance_ = 1000000-70
    vertical_tranlation_between_projection_ = 0
    if not os.path.exists(os.path.join(cur_dir_, "output")):
        os.mkdir(os.path.join(cur_dir_, "output"))
    tmp_dir = os.path.join(cur_dir_, "output", model_name_)
    if not os.path.exists(tmp_dir):
        os.mkdir(tmp_dir)
    tmp_dir = os.path.join(tmp_dir, "{}_{}".format(model_name_, scale_))
    if not os.path.exists(tmp_dir):
        os.mkdir(tmp_dir)
    dose_save_path = os.path.join(tmp_dir, "{}_{}_dose.dat".format(model_name_, scale_))
    projection_save_path = os.path.join(tmp_dir, "{}_{}".format(model_name_, scale_))

    print("projection will be saved in {}".format(projection_save_path))
    # projection_save_path = os.path.join(cur_dir_, "output", "{}/{}_{}".format(model_name_,
    #                                                                           model_name_, scale_),
    #                                     "{}_{}".format(model_name_, model_name_, scale_))
    # software will add .raw at the end autonomously

    # voxelized geometry setting todo
    if vox_file_path_ is None:
        vox_file_path_ = os.path.join(cur_dir_, "vox", "{}/{}_{}_{}.vox".format(model_name_, model_name_,
                                                                                float(scale_),
                                                                                time_stamp_))
    else:
        vox_file_path_ = os.path.join("vox", vox_file_path_)

    print("to read vox file path is {}".format(vox_file_path_))

    # material setting, please notice the order todo
    mat_list_ = ["air__5-120keV.mcgpu.gz",
                 "adipose_ICRP110__5-120keV.mcgpu.gz",
                 "brain_ICRP110__5-120keV.mcgpu.gz",
                 "bone_ICRP110__5-120keV.mcgpu.gz",
                 "blood_ICRP110__5-120keV.mcgpu.gz",]

    if not os.path.exists(kvp_path_):
        if kvp_path_.split(".")[-1] != "spc":
            raise RuntimeError("please notice that the file suffix is .spc,"
                               "while the input file is {}".format(kvp_path_))
        else:
            raise RuntimeError("please check the path the of the input file,"
                               "current input path is {}".format(kvp_path_))

    if not os.path.exists(vox_file_path_):
        raise RuntimeError("please check the input .vox file path, and "
                           "current input vox path is {}".format(vox_file_path_))

    # write in the .in file
    print("start to write the .in file")
    print(".in file will be saved in {}".format(save_in_file_path_))
    file_ = open(save_in_file_path_, "w")
    file_.write('# >>>> INPUT FILE FOR MC-GPU v1.3 >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>\n')
    file_.write("this file was created in".format(add_timestamp()))
    file_.write('\n')

    # [SECTION SIMULATION CONFIG v.2009-05-12]
    file_.write('# [SECTION SIMULATION CONFIG v.2009-05-12]\n')
    file_.write(str(photon_num_) + '\n')
    file_.write(str(random_seed_) + "\n")  # rand seed
    file_.write(str(gpu_no_) + '\n')
    file_.write(str(thread_num_) + '\n')
    file_.write(str(simulated_num_per_thread_) + '\n')
    file_.write('\n')

    # [SECTION SOURCE v.2011-07-12]
    file_.write('# [SECTION SOURCE v.2011-07-12]\n')
    file_.write(kvp_path_ + '\n')
    file_.write(str(source_position_x_) + ' ' + str(source_position_y_) + ' ' + str(source_position_z_) + '\n')
    file_.write('0 1 0\n')
    file_.write('-1 -1\n')
    file_.write('\n')

    # [SECTION IMAGE DETECTOR v.2009-12-02]
    file_.write('# [SECTION IMAGE DETECTOR v.2009-12-02]\n')
    file_.write('1\n')  # output format - raw data
    file_.write(projection_save_path + '\n')
    file_.write(str(nu_) + ' ' + str(nv_) + '\n')
    file_.write(str(nu_ * du_) + ' ' + str(nv_ * dv_) + '\n')
    file_.write(str(sdd_) + '\n')
    file_.write('\n')

    # [SECTION CT SCAN TRAJECTORY v.2011-10-25]
    file_.write('# [SECTION CT SCAN TRAJECTORY v.2011-10-25]\n')
    file_.write(str(projection_num_) + '\n')
    file_.write(str(angle_between_projection_) + '\n')
    file_.write('{}  {}\n'.format(angle_of_interest_[0], angle_of_interest_[1]))
    file_.write(str(source_to_ratation_axis_distance_) + '\n')
    file_.write('0.0\n')  # for helical
    file_.write('\n')

    if not cal_dose_:
        # [SECTION DOSE DEPOSITION v.2012-12-12]
        file_.write('# [SECTION DOSE DEPOSITION v.2012-12-12]\n')
        file_.write('NO\n')
        file_.write('NO\n')
        file_.write('dose\n')
        file_.write('1 3\n')
        file_.write('1 1\n')
        file_.write('1 2\n')
    else:
        # [SECTION DOSE DEPOSITION v.2012-12-12]
        file_.write('\n')
        file_.write('# [SECTION DOSE DEPOSITION v.2012-12-12]\n')
        file_.write('YES\n')
        file_.write('YES\n')
        file_.write(dose_save_path+"\n")
        file_.write('1 {}\n'.format(nx))
        file_.write('1 {}\n'.format(ny))
        file_.write('1 {}\n'.format(nz))
        file_.write('\n')


# [SECTION VOXELIZED GEOMETRY FILE v.2009-11-30]
    file_.write('# [SECTION VOXELIZED GEOMETRY FILE v.2009-11-30]\n')
    file_.write(vox_file_path_ + '\n')
    file_.write('\n')

    # [SECTION MATERIAL FILE LIST v.2009-11-30]
    file_.write('# [SECTION MATERIAL FILE LIST v.2009-11-30]\n')
    for mat in mat_list_:
        file_.write(os.path.join(cur_dir_, "material", mat + "\n"))
    file_.write('\n')

    file_.write('# >>>> END INPUT FILE >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>')
    print(".in file is done")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="define the model name")
    parser.add_argument("-name", "--model_name", default="model1109", type=str)
    parser.add_argument("-scale", "--scale", default="0.7", type=str)
    parser.add_argument("-vox", "--vox_path",
                        default=r"/mnt/520/dcp/demo/vox/medical_2/medical_2_0.045_.vox",
                        type=str)
    parser.add_argument("-pho_num", "--photon_num", default="3082003700")
    parser.add_argument("-gpu", "--gpu", default="3")
    parser.add_argument("-dose", "--dose", default=True)
    parser.add_argument("-dose_scale", "--dose_scale", default="1")

    args = parser.parse_args()
    print(args)

    # time_stamp = add_timestamp()
    time_stamp = ""
    print(args.vox_path)
    # create_in_file(scale_=0.9, model_name_="model1109", time_stamp_=time_stamp)
    scale_tmp = float(args.scale)
    vox_file_path = args.vox_path.strip("\r")
    create_in_file(scale_=scale_tmp, model_name_=args.model_name,
                   time_stamp_=time_stamp,
                   vox_file_path_=vox_file_path,
                   photon_num_=args.photon_num,
                   gpu_=args.gpu,
                   cal_dose_=args.dose,
                   dose_scale_=args.dose_scale)
